# import argparse

# def parse_args():
#     """解析命令行参数"""
#     parser = argparse.ArgumentParser(description='示例脚本：接收CUDA设备和端口参数')
    
#     # 添加参数
#     parser.add_argument(
#         '-c',
#         '--cuda_id', 
#         type=str,
#         required=True,  # 必须传入
#         help='CUDA设备ID，例如 "0" 或 "0,1"（字符串类型）'
#     )
    
#     # 解析参数
#     args = parser.parse_args()
#     return args
# args = parse_args()

import os
os.environ['CUDA_VISIBLE_DEVICES']='1'

from utils import util_for_huggingface

import torch,pdb
from diffusers import FluxControlPipeline, FluxPriorReduxPipeline
from diffusers.utils import load_image
from util_flux import process_img_1024,vertical_concat_images,horizontal_concat_images
from util_flux import resize_with_aspect
from PIL import Image,ImageOps
from image_gen_aux import DepthPreprocessor
from itertools import product
from util_for_os import osj,ose
import numpy as np


from MODEL_CKP import DEPTH_PREDCITION,FLUX_REDUX,FLUX_DEPTH,FLUX

# types = ['collar','sleeve','pockets']
# examples_dir = lambda t: f'/mnt/nas/shengjie/datasets/cloth_{t}_balanced'
output_dir = 'tmp_redux0901'
os.makedirs(output_dir,exist_ok=True)

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                    FLUX_REDUX, 
                                    torch_dtype=torch.bfloat16).to("cuda")
depth_processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)

pipe = FluxControlPipeline.from_pretrained(FLUX_DEPTH, torch_dtype=torch.bfloat16).to("cuda")


from PIL import ImageFilter

from demo_rmbg import load_rmbg,get_mask_by_rmbg
model_rmbg = load_rmbg()


# 各取五张 然后交叉 product 生成
num = 5

def get_tmp_clo1_clo2(dir):
    # dir = './compare_longshort'
    files = os.listdir(dir)
    files.sort()
    files = [ osj( dir , f ) for f in files]
    return files[:5],files[-5:]

process_dirs = [
    'compare_longshort',
    'compare_collar',
    'compare_pockets',
    'compare_sleeve',
]

from utils.util_for_depthcontrol import get_result


for process_dir in process_dirs:

    assert os.path.exists( process_dir ),process_dir

for process_dir in process_dirs:
    
    clo_list1,clo_list2 = get_tmp_clo1_clo2(process_dir) 
    for process_type in ['_normal','_reverse']:
        clo_list1,clo_list2 = (clo_list1,clo_list2) if process_type == '_normal' else (clo_list2,clo_list1)

        for gaussian_score in [10,20,30,40,50]:

            def get_mask(size=(1024,1024),fill=(255,255,255)):
                return Image.new('RGB',size,color=fill)
            res_imgs = [
                [get_mask()]+[ process_img_1024(path) for path in clo_list1 ],
                *[[process_img_1024(path)] for path in clo_list2], 
            ]
            idx = 0
            for c1,c2 in product(clo_list1,clo_list2):
                '''
                构建
                /        原图 c1-1  c1-2  c1-3 ...
                原图c2-1     合成1  合成2 合成3 ...
                c2-2        ...
                # c2-3
                ...
                '''

                control_image_path = c1
                redux_image_path = c2


                control_image = process_img_1024( control_image_path )

                redux_image = process_img_1024( redux_image_path )

                alpha = 0.8 
                control_alpha = 0.4
                image = get_result( 
                    control_image , redux_image,
                    depth_processor = depth_processor,
                    redux_pipe = pipe_prior_redux, flux_pipe = pipe,
                    get_mask = get_mask_by_rmbg, mask_model = model_rmbg,
                    alpha = alpha ,control_alpha = control_alpha,gaussian_score=gaussian_score 
                )

                res_imgs[idx%num + 1].append(
                    process_img_1024('',img_pil=image)
                )

                idx += 1 
                print('cur id ',idx)


                # concat_tmp_res = horizontal_concat_images( [ 
                #     process_img_1024('',img_pil=control_image) , 
                #     process_img_1024(control_image_path),
                #     redux_image ,   
                #     process_img_1024('',img_pil=image) ,
                #     ] )

                # concat_tmp_res.save('tmp_depth_control.jpg')

            res_imgs_hori_concat = [
                horizontal_concat_images(ri) for ri in res_imgs
            ]
            res_imgs_verti_concat = vertical_concat_images(res_imgs_hori_concat)
            res_imgs_verti_concat.save(f'{output_dir}/{process_dir}{process_type}_{gaussian_score}.jpg')